import cv2
import numpy as np
from scipy import signal
from skimage.draw import line_aa

from helpers import *


bbox = [[514, 276, 586, 308], [619, 289, 686, 320], [1097, 252, 1134, 386]]
img_fns = ['T_20220313134425_400.jpg', 'T_20220313134425_440.jpg', 'T_20220313134426_160.jpg']

img1 = cv2.imread(img_fns[0])
img2 = cv2.imread(img_fns[1])
img3 = cv2.imread(img_fns[2])

img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY).astype(np.float) / 255
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY).astype(np.float) / 255
img3 = cv2.cvtColor(img3, cv2.COLOR_BGR2GRAY).astype(np.float) / 255

img_mean = (img1 + img2) / 2

imgsub = np.abs(img1 - img2)

fg_roi = img1[bbox[0][1]: bbox[0][3], bbox[0][0]: bbox[0][2]].copy()
diam = round(max(fg_roi.shape[0:2]))
r = int(diam / 2)
cy, cx = int((bbox[0][1] + bbox[0][3]) / 2), int((bbox[0][0] + bbox[0][2]) / 2)
F = img1[cy - r: cy + r, cx - r: cx + r].copy()

img1[bbox[0][1]: bbox[0][3], bbox[0][0]: bbox[0][2]] = img2[bbox[0][1]: bbox[0][3], bbox[0][0]: bbox[0][2]]
bg_img = img1

# cv2.rectangle(img1, (bbox[0][0], bbox[0][1]), (bbox[0][2], bbox[0][3]), (0, 0, 255), 1, cv2.LINE_AA)
# cv2.rectangle(img1, (bbox[1][0], bbox[1][1]), (bbox[1][2], bbox[1][3]), (0, 0, 255), 1, cv2.LINE_AA)
# cv2.rectangle(img2, (bbox[0][0], bbox[0][1]), (bbox[0][2], bbox[0][3]), (0, 0, 255), 1, cv2.LINE_AA)
# cv2.rectangle(img2, (bbox[1][0], bbox[1][1]), (bbox[1][2], bbox[1][3]), (0, 0, 255), 1, cv2.LINE_AA)

H = np.zeros(bg_img.shape[0:2])
rr, cc, val = line_aa(int((bbox[0][1] + bbox[0][3]) / 2), int((bbox[0][0] + bbox[0][2]) / 2),
                      int((bbox[1][1] + bbox[1][3]) / 2), int((bbox[1][0] + bbox[1][2]) / 2))
H[rr, cc] = val

# cv2.imshow('H', H)

ratio = 0.3
B = img2 * ratio + bg_img * (1 - ratio)
rad = diam / 2
M = diskMask(rad)
FM = F * M

H = H / np.sum(H)

HM = signal.fftconvolve(H, M, mode='same')
HF = np.zeros(B.shape)
HF = signal.fftconvolve(H, FM, mode='same')

im = B * (1 - HM) + HF

cv2.imshow('HM', HM)

im[bbox[2][1]: bbox[2][3], bbox[2][0]: bbox[2][2]] = img3[bbox[2][1]: bbox[2][3], bbox[2][0]: bbox[2][2]]

cv2.imshow('img1', bg_img)
cv2.imshow('img2', im)
cv2.imshow('imgsub', imgsub)

hm1s = ((np.max(HM) - HM) * 255).astype(np.uint8)
hm1s[hm1s == 178] = 255

cv2.imwrite('T_20220313134425_result.jpg', (im * 255).astype(np.uint8))
cv2.imwrite('T_20220313134425_hm.jpg', (HM * 255).astype(np.uint8))
cv2.imwrite('T_20220313134425_hm1s.jpg', hm1s)

cv2.waitKey(2000)
